#include "huffman.hpp"
#include <iostream>
#include <algorithm>
#include <map>

namespace Huffman{
  // print out the frequency table
  // debug utility
  void print_frequencies(std::map<const int, int>& ftable){
    std::cout << "------------------------------- " << std::endl;
    std::cout << "Symbol frequencies: " << std::endl;
    std::cout << "------------------------------- " << std::endl;
    std::map<const int, int>::iterator iter;
    for(iter = ftable.begin(); iter != ftable.end(); iter++){
      std::cout << (*iter).first << " | " << (*iter).second << std::endl;
    }
    std::cout << "------------------------------- " << std::endl;
  }

  // print out the Huffman code for each message symbol
  // debug utility
  void print_codes(std::map<const int, std::vector<int> > &lookupTable){
    std::cout << "------------------------------- " << std::endl;
    std::cout << "Huffman codes: " << std::endl;
    std::cout << "------------------------------- " << std::endl;
    std::map<const int, std::vector<int> >::iterator iter;
    for(iter = lookupTable.begin(); iter != lookupTable.end(); iter++){
      std::cout << (*iter).first << " | ";
      for(int i = 0; i < (signed int) (*iter).second.size(); ++i)
	std::cout << ((*iter).second)[i];
      std::cout << std::endl;
    }
    std::cout << "------------------------------- " << std::endl;
  }

  // print out the Huffman code for each message symbol
  // debug utility
  void print_codes_GF2(std::map<const int, vec_GF2 > &lookupTable){
    std::cout << "------------------------------- " << std::endl;
    std::cout << "Huffman codes: " << std::endl;
    std::cout << "------------------------------- " << std::endl;
    std::map<const int, vec_GF2 >::iterator iter;
    for(iter = lookupTable.begin(); iter != lookupTable.end(); iter++){
      std::cout << (*iter).first << " | " << (*iter).second << std::endl;
    }
    std::cout << "------------------------------- " << std::endl;
  }

  // computes the frequency of each symbol in a vector
  // FIXME: to be rewritten using template
  void compute_frequency(std::vector<int> symbols, 
			 std::map<const int, int> &ftable){
    std::sort(symbols.begin(), symbols.end());
    int lastelement = 0;
    for(long i = 0; i < (signed long) symbols.size(); ++i){
      if (i == 0){
	lastelement = symbols[i];
	ftable[symbols[i]] = 1;
      }
      else{
	if (lastelement != symbols[i]){
	  ftable[symbols[i]] = 1;
	  lastelement = symbols[i];
	}
	else
	  ftable[symbols[i]] ++;
      }
    }
  }

  // create a tree node for each entry in the frequency table
  void initialize(std::vector<TreeNode*>& trees, 
		  std::map<const int, int> &ftable){
    std::map<const int, int>::iterator iter;
    for(iter = ftable.begin(); iter != ftable.end(); iter++){
      TreeNode* tree = new TreeNode();
      tree -> symbol = (*iter).first;
      tree -> weight = (*iter).second;
      trees.push_back(tree);
    }
  }

  // sort the trees by weight, and merge the two whose weights are
  // minimum
  void sort_and_merge(std::vector<TreeNode*>& trees){
    std::sort(trees.begin(), trees.end(), lt());
    if (trees.size() >= 2) {
      TreeNode* tree = new TreeNode();
      tree -> weight = trees[0] -> weight + trees[1] -> weight;
      tree -> leftChild = trees[0];
      tree -> rightChild = trees[1];
      trees.erase (trees.begin(), trees.begin()+2);
      trees.push_back(tree);
      sort_and_merge(trees);
    }
    else
      return;
  }

  // create a Huffman tree
  void build_huffman_tree(std::vector<TreeNode*> &trees, 
			  std::map<const int, int> &ftable){
    trees.clear(); // may have memory leaking
    initialize(trees, ftable);
    sort_and_merge(trees);
  }

  // determine the codeword for each symbol by traversing the Huffman
  // tree
  void traverse(TreeNode* tree, std::vector<int> &codewordStack, 
		std::map<const int, std::vector<int> > &lookupTable)
  {
    if (tree -> is_leaf()){
      const int key = tree -> symbol;
      std::vector<int> val (codewordStack.begin(), codewordStack.end());  
      lookupTable[key] = val;
    }
    else{
      if (tree -> leftChild){
	codewordStack.push_back(0);
	traverse(tree -> leftChild, codewordStack, lookupTable);
	codewordStack.pop_back();
      }
      if (tree -> rightChild){
	codewordStack.push_back(1);
	traverse(tree -> rightChild, codewordStack, lookupTable);
	codewordStack.pop_back();
      }
    }
  }

  // convert a bitVector with type std::vector to a vector with type
  // vec_GF2
  void to_vec_GF2(vec_GF2& result, const std::vector<int> &bitVector){
    result.SetLength(bitVector.size());
    for(long i = 0; i < (signed long) bitVector.size(); ++i){
      result(i+1) = to_GF2(bitVector[i]);
    }
  }

  // determine the codeword for each symbol by traversing the Huffman
  // tree
  void traverse_GF2(TreeNode* tree, std::vector<int> &codewordStack, 
		    std::map<const int, vec_GF2 > &lookupTable)
  {
    if (tree -> is_leaf()){
      const int key = tree -> symbol;
      //std::vector<int> val (codewordStack.begin(), codewordStack.end());
      vec_GF2 codesymbol;
      to_vec_GF2(codesymbol, codewordStack);
      lookupTable[key] = codesymbol;
    }
    else{
      if (tree -> leftChild){
	codewordStack.push_back(0);
	traverse_GF2(tree -> leftChild, codewordStack, lookupTable);
	codewordStack.pop_back();
      }
      if (tree -> rightChild){
	codewordStack.push_back(1);
	traverse_GF2(tree -> rightChild, codewordStack, lookupTable);
	codewordStack.pop_back();
      }
    }
  }


  // compute the codewords for each symbol
  void get_codewords(std::vector<TreeNode*> &trees, 
		     std::vector<int> &codewordStack, 
		     std::map<const int, std::vector<int> > &lookupTable){
    codewordStack.clear();
    traverse(trees[0], codewordStack, lookupTable);
  }

  // compute the codewords for each symbol
  void get_codewords_GF2(std::vector<TreeNode*> &trees, 
			 std::vector<int> &codewordStack, 
			 std::map<const int, vec_GF2 > &lookupTable){
    codewordStack.clear();
    traverse_GF2(trees[0], codewordStack, lookupTable);
  }


  // computing the Huffman codes of all the symbols
  void compute_lookup_table(std::vector<int>& symbols, 
			    std::map<const int, vec_GF2 >& lookupTable, 
			    std::vector<TreeNode*>& trees){
    std::map<const int, int> ftable;
    std::vector<int> codewordStack;
    compute_frequency(symbols, ftable);
    build_huffman_tree(trees, ftable);
    get_codewords_GF2(trees, codewordStack, lookupTable);

#ifdef VERBOSE
    print_frequencies(ftable);
    print_codes_GF2(lookupTable);
#endif

  }


  // the encoder of Huffman codes
  void huffman_encode(vec_GF2& codeword,
		      std::vector<int>& symbols, 
		      std::vector<TreeNode*>& trees){
    std::map<const int, vec_GF2 > lookupTable;
    compute_lookup_table(symbols, lookupTable, trees);
    for(long i = 0; i < (signed long) symbols.size(); ++i)
      append(codeword, lookupTable[symbols[i]]);
  }





  // the decoder of Huffman codes
  void huffman_decode(std::vector<int>& symbols, 
		      vec_GF2& codeword, 
		      std::vector<TreeNode*>& trees) {
    TreeNode* tree = trees[0];
    TreeNode* ptr = tree;
    for(long i = 1; i <= codeword.length(); ++i){
      if (codeword(i) == 1) {
	if (!(ptr -> rightChild)){
	  std::cout << "Warning: Input codeword contains error" << std::endl;
	  symbols.clear();
	  return;
	}
	ptr = ptr -> rightChild;
      }
      else{
	if (!(ptr -> leftChild)){
	  std::cout << "Warning: Input codeword contains error" << std::endl;
	  symbols.clear();
	  return;
	}
	ptr = ptr -> leftChild;
      }
      if (ptr -> is_leaf()) {
	symbols.push_back(ptr->symbol);
	ptr = tree;  // decoded one symbol successfully, go back to the root.
      }
    }
  }

} // Huffman
